// Copyright (c) 2020-present,  INSPUR Co, Ltd.  All rights reserved.
// This source code is licensed under Apache 2.0 License.

#include "db/db_test_util.h"
#include "memory_arena_art_index.h"
#include "pure_mem/pmemrep.h"
#include "pure_mem/rangearena/range_arena.h"
#include "util/random.h"
#include "util/testharness.h"

namespace rocksdb{
    class TestKeyComparator : public MemTableRep::KeyComparator {
    public:
        typedef rocksdb::Slice DecodedType;

        DecodedType decode_key(const char* key) const override {
          // The format of key is frozen and can be terated as a part of the API
          // contract. Refer to MemTable::Add for details.
          return GetLengthPrefixedSlice(key);
        }

        // Compare a and b. Return a negative value if a is less than b, 0 if they
        // are equal, and a positive value if a is greater than b
        // Compare size in byte order
        int operator()(const char* prefix_len_key1,
                       const char* prefix_len_key2) const override {
          Slice k1 = GetLengthPrefixedSlice(prefix_len_key1);
          Slice k2 = GetLengthPrefixedSlice(prefix_len_key2);

          int comLength = std::min(k1.size(), k2.size());
          for (int i = 0; i < comLength; i++) {
            if ((uint8_t)k1[i] > (uint8_t)k2[i]) {
              return 1;
            } else if ((uint8_t)k1[i] < (uint8_t)k2[i])
              return -1;
          }

          if (k1.size() > k2.size()) return 1;
          if (k1.size() < k2.size()) return -1;
          return 0;
        }

        int operator()(const char* prefix_len_key, const Slice& key) const override {
          Slice key_pre = GetLengthPrefixedSlice(prefix_len_key);
          return key_pre.compare(key);
        }

        ~TestKeyComparator() override = default;
    };

    typedef uint64_t Key;
    typedef PureMemRep TestMemArenaRep;
    const size_t standard_key = 1000000;

    static const char* Encode(const Key* key) {
      std::string str = std::to_string(*key);
      int len = (int)(str.length()) + 8 + 1;
      uint32_t keySize = VarintLength(len) + len;
      char* ret = new char[keySize];
      memset(ret , '\0', keySize);
      char*p = EncodeVarint32(ret, len);
      memcpy(p, str.c_str(), str.length());
      return reinterpret_cast<const char*>(ret);
    }

    static Slice Decode(const char* key) {
      Slice buf = GetLengthPrefixedSlice(key);
      Slice mvccKey = Slice(buf.data(), buf.size() - 8);
      const char ts_size = mvccKey[mvccKey.size() - 1];
      if ((size_t)ts_size >= mvccKey.size()) {
        return nullptr;
      }
      return {buf.data(), mvccKey.size() - 1 - ts_size};
    }


    class MemArenaARTIndexTest : public testing::Test {
    public:
        static void Insert(TestMemArenaRep* list, Key key) {
          char** buf = nullptr;
          void* rangeArena;
          const auto& keyy = (const Slice&) key;
          const size_t len = sizeof(Key);
          *buf = (char*)list->AllocatePure(len, buf, keyy, &rangeArena);
          memcpy(*buf, &key, sizeof(Key));
          void* buff = *buf;
          list->Insert(buff);
        }

        void Validate(TestMemArenaRep* list) {
          // Check keys exist.
          for (Key key : keys_) {
            ASSERT_TRUE(list->Contains(Encode(&key)));
          }
          // Iterate over the list, make sure keys appears in order and no extra
          // keys exist.
          InlineUserKeyIndex<const MemTableRep::KeyComparator&>::Iterator iter =
              (InlineUserKeyIndex<const MemTableRep::KeyComparator&>::Iterator)
                  reinterpret_cast<const InlineUserKeyIndex<const MemTableRep::KeyComparator &> *>(list->GetIterator());
          ASSERT_FALSE(iter.Valid());
          Key zero = 0;
          iter.Seek(Encode(&zero));
          for (Key key : keys_) {
            ASSERT_TRUE(iter.Valid());
            ASSERT_EQ((char*)key, Decode(iter.key()).data());
            iter.Next();
          }
          ASSERT_FALSE(iter.Valid());
        }

    private:
        std::set<Key> keys_;
    };

    Slice genMVCCKey1(const char* key) {
      uint32_t keySize = strlen(key) + 1;
      char* ret = new char[keySize];
      memset(ret, '\0', keySize);
      memcpy(ret, key, strlen(key));
      return Slice(ret, keySize);
    }


    Slice GenInternalKey(const Slice &key, char *buf){
      size_t key_size = key.size() + 1;
      size_t internal_key_size = key_size + 8;

      char* p = buf;
      memset(p, '\0', key_size);
      memcpy(p, key.data(), key.size());
      p += key_size;
      memset(p, 255, 8);
      return Slice(buf, internal_key_size);

    }

    //测试第一层索引为空的Seek
    TEST_F(MemArenaARTIndexTest, SeekForEmptyAndTest) {
      std::cout<<"SeekForEmptyAndTest  "<<std::endl;
      std::set<Key> keys;
      InternalKeyComparator comparator(BytewiseComparator());
      Options options = Options();
      ImmutableCFOptions ioptions(options);
      const TestKeyComparator compare;

      int bound = 1;
      auto manager = new MemArenaARTIndex(comparator);
      MemArena *memoryArenas[4];
      std::string us[4];
      Slice userkeys[4];
      char* buf[4];
      size_t encode_len = 0;

      //将startkey和MemoryArena插入第一层ART索引
      for (int i = 0; i < 4; i++) {
        memoryArenas[i] = new MemArena(comparator, ioptions);
        us[i] = std::to_string(bound + 2 * i);
        userkeys[i] = Slice(us[i].data(), us[i].size());

        encode_len = us[i].size() + 1 + 8;
        buf[i] = new char[encode_len];
        Slice internal_key = Slice(GenInternalKey(userkeys[i], buf[i]));
        auto memory = manager->SeekForMemoryArena(internal_key);
        if(memory != nullptr){
          std::cout<<"查到的memoryarena为："<<(void*)memory<<std::endl;
          exit(1);
        }
      }
      delete manager;
      for(int i = 0; i < 4; i++){
        delete[] buf[i];
        delete memoryArenas[i];
      }
    }

    //测试第一层索引的插入及查询
    TEST_F(MemArenaARTIndexTest, InsertAndTest) {
      std::set<Key> keys;
      InternalKeyComparator comparator(BytewiseComparator());
      Options options = Options();
      ImmutableCFOptions ioptions(options);
      const TestKeyComparator compare;

      int bound = 1;
      auto manager = new MemArenaARTIndex(comparator);
      MemArena *memoryArenas[4];
      std::string us[4];
      Slice userkeys[4];
      char* buf[4];
      size_t encode_len = 0;

      //测试时第一个单独处理
      char start_key[2];
      start_key[0] = 0;
      start_key[1] = 0;
      Slice start(start_key, 2);
      MemArena memoryArena(comparator, ioptions);
      manager->Insert(start, &memoryArena);

      //将startkey和MemoryArena插入第一层ART索引
      for (int i = 0; i < 4; i++) {
        memoryArenas[i] = new MemArena(comparator, ioptions);
        us[i] = std::to_string(bound + 2 * i);
        userkeys[i] = Slice(us[i].data(), us[i].size());

        bool ok = manager->Insert(userkeys[i], memoryArenas[i]);
        ASSERT_TRUE(ok);

        encode_len = us[i].size() + 1 + 8;
        buf[i] = new char[encode_len];
        Slice internal_key = Slice(GenInternalKey(userkeys[i], buf[i]));
        auto memory = manager->SeekForMemoryArena(internal_key);
        if(memoryArenas[i] != memory){
          std::cout<<us[i]<<"对应插入的memoryarena :"<<(void*)memoryArenas[i]<<std::endl;
          std::cout<<"查到的memoryarena为："<<(void*)memory<<std::endl;
        }
      }
      delete manager;
      for(int i = 0; i < 4; i++){
        delete[] buf[i];
        delete memoryArenas[i];
      }
    }

    //测试第一层索引的插入及查询，及相同startKey时的处理
    TEST_F(MemArenaARTIndexTest, InsertSameKeyAndTest) {
      std::set<Key> keys;
      InternalKeyComparator comparator(BytewiseComparator());
      Options options = Options();
      ImmutableCFOptions ioptions(options);
      const TestKeyComparator compare;

      int bound = 1;
      auto manager = new MemArenaARTIndex(comparator);
      MemArena *memoryArenas[6];
      std::string us[6];
      Slice userkeys[6];
      char* buf[6];
      size_t encode_len = 0;

      //测试时第一个单独处理
      char start_key[2];
      start_key[0] = 0;
      start_key[1] = 0;
      Slice start(start_key, 2);
      MemArena memoryArena(comparator, ioptions);
      manager->Insert(start, &memoryArena);

      //将startkey和MemoryArena插入第一层ART索引
      for (int i = 0; i < 6; i++) {
        memoryArenas[i] = new MemArena(comparator, ioptions);
        us[i] = std::to_string(bound + 2 * i);
        if(i == 5){
          us[i] = std::to_string(bound + 4);
        }
        userkeys[i] = Slice(us[i].data(), us[i].size());
        std::cout<<userkeys[i].ToString()<<"------"<<memoryArenas[i]<<std::endl;
        bool ok = manager->Insert(userkeys[i], memoryArenas[i]);
        ASSERT_TRUE(ok);

        encode_len = us[i].size() + 1 + 8;
        buf[i] = new char[encode_len];
        Slice internal_key = Slice(GenInternalKey(userkeys[i], buf[i]));
        auto memory = manager->SeekForMemoryArena(internal_key);
        if(memoryArenas[i] != memory){
          std::cout<<us[i]<<"对应插入的memoryarena :"<<(void*)memoryArenas[i]<<std::endl;
          std::cout<<"查到的memoryarena为："<<(void*)memory<<std::endl;
        }
      }
      delete manager;
      for(int i = 0; i < 6; i++){
        delete[] buf[i];
        delete memoryArenas[i];
      }
    }


    //测试数据插入时Seek到的MemoryArena是否正确，
    //向第二层ART索引插入kv数据，首先在第一层找到应该插入哪个MemoryArena
    TEST_F(MemArenaARTIndexTest, InsertAndSeekTest) {
      std::set<Key> keys;
      InternalKeyComparator comparator(BytewiseComparator());
      Options options = Options();
      ImmutableCFOptions ioptions(options);
      const TestKeyComparator compare;
      char tmp[1000] = {'1'};
      Slice value(tmp, 1000);
      SequenceNumber s = 100000;
      ValueType type = kTypeValue;
      int bound = 1;
      auto manager = new MemArenaARTIndex(comparator);
      MemArena *memoryArenas[4];
      std::string us[4];
      Slice userkeys[4];
      int mul = 0;

      //测试时第一个单独处理
      char start_key[2];
      start_key[0] = 0;
      start_key[1] = 0;
      Slice start(start_key, 2);
      MemArena memoryArena(comparator, ioptions);
      manager->Insert(start, &memoryArena);
      //将startkey和MemoryArena插入第一层ART索引
      for (int i = 0; i < 4; i++) {
        memoryArenas[i] = new MemArena(comparator, ioptions);
        us[i] = std::to_string(bound + 2 * i);
        userkeys[i] = Slice(us[i].data(), us[i].size());
        manager->Insert(userkeys[i], memoryArenas[i]);
    }

      //向第二层ART索引插入kv数据
      for (int i = 1; i < 10000; i++) {
        mul = i / 1000;
        s++;
        std::string a = std::to_string(mul) + std::to_string(i) + "_test";

        Slice key = Slice(genMVCCKey1(a.data()));

        size_t key_size = key.size();
        size_t val_size = value.size();
        size_t internal_key_size = key_size + 8;

        char buf[internal_key_size];
        char* p = buf;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed = PackSequenceAndType(s, type);
        EncodeFixed64(p, packed);
        char ret[val_size];
        memcpy(ret, value.data(), val_size);

        Slice slice_key(buf, internal_key_size);
        Slice slice_value(ret, val_size);

        MemArena* addr = dynamic_cast<MemArena *>(manager->SeekForMemoryArena(slice_key));
        if(addr == nullptr){
          std::cout<<"没有找到key对应得memoryarena, key:"<<key.ToString()<<std::endl;
          exit(1);
        }
        switch (mul){
          case 0:
            if(addr != &memoryArena)
              std::cout<<"原MemoryArena ："<<(&memoryArena)<<",查到的MemoryArena ： "<<addr<<std::endl;
            break;
          case 1:
          case 2:
            if(addr != memoryArenas[0])
              std::cout<<"原MemoryArena ："<<(memoryArenas[0])<<",查到的MemoryArena ： "<<addr<<std::endl;
            break;
          case 3:
          case 4:
            if(addr != memoryArenas[1])
              std::cout<<"原MemoryArena ："<<(memoryArenas[1])<<",查到的MemoryArena ： "<<addr<<std::endl;
            break;
          case 5:
          case 6:
            if(addr != memoryArenas[2])
              std::cout<<"原MemoryArena ："<<(memoryArenas[2])<<",查到的MemoryArena ： "<<addr<<std::endl;
            break;
          default:
            if(addr != memoryArenas[3])
              std::cout<<"原MemoryArena ："<<(memoryArenas[3])<<",查到的MemoryArena ： "<<addr<<std::endl;
            break;
        }

        bool res = addr->KVDataInsert(slice_key, slice_value);
        ASSERT_TRUE(res);
      }
      delete manager;
      for(auto & memoryarena : memoryArenas){
        delete memoryarena;
      }
    }


//    void insertAndSeek(size_t min, size_t max, MemArenaARTIndex* memIndex){
//      std::string u;
//      Slice userkey;
//      InternalKeyComparator comparator(BytewiseComparator());
//      Options options = Options();
//      ImmutableCFOptions ioptions(options);
//      for(int i = min; i < max; i++){
//        u = std::to_string(i);
//        userkey = Slice(u.data(), u.size());
//        auto memoryArena = new MemArena(comparator, ioptions);
//        memIndex->Insert(userkey, memoryArena);
//      }
//    }
//    /////多线程测试------待完成
//    TEST_F(MemArenaARTIndexTest, MultiInsertAndSeekTest) {
//      std::set<Key> keys;
//      InternalKeyComparator comparator(BytewiseComparator());
//      Options options = Options();
//      ImmutableCFOptions ioptions(options);
//      const TestKeyComparator compare;
//      char tmp[1000] = {'1'};
//      Slice value(tmp, 1000);
//      SequenceNumber s = 100000;
//      ValueType type = kTypeValue;
//      int bound = 1;
//      auto manager = new MemArenaARTIndex(comparator, ioptions);
//      MemArena *memoryArenas[4];
//      std::string us[4];
//      Slice userkeys[4];
//
//      //测试时第一个单独处理
//      char start_key[2];
//      start_key[0] = 0;
//      start_key[1] = 0;
//      Slice start(start_key, 2);
//      MemArena memoryArena(comparator, ioptions);
//      manager->Insert(start, &memoryArena);
//      //将startkey和MemoryArena插入第一层ART索引
//
//      std::thread test1(insertAndSeek, 1, 3, manager);
//      std::thread test2(insertAndSeek, 2, 4, manager);
//      std::thread test3(insertAndSeek, 4, 6, manager);
//      std::thread test4(insertAndSeek, 6, 9, manager);
//      test1.join();
//      test2.join();
//      test3.join();
//      test4.join();
//
//      delete manager;
//    }


}
int main(int argc, char** argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}